import copy
import os


class BaseExporter(object):
    def __init__(self, export_dir="./data/", export_type="csv", header=True):
        self.base_config = {
            "export_dir": export_dir,
            "export_type": export_type,
            "header": header,
            "file_name": None
        }

    def export_table_param(self):
        raise NotImplementedError()


class CSVExport(BaseExporter):
    def __init__(self, export_dir, header=True):
        super().__init__(export_dir, export_type="csv", header=header)

    def export_table_param(self, export_config, table_records: dict, tab_cols=None):
        assert export_config is not None
        assert export_config["file_name"] is not None
        if not os.path.exists(export_config["export_dir"]):
            os.makedirs(export_config["export_dir"])

        if export_config["header"] and tab_cols is None:
            print("[ERROR] header is True, but tab_cols is None")
            return

        file_path = os.path.join(export_config["export_dir"], export_config["file_name"])
        if export_config["export_type"] == "csv":
            import csv
            with open(file_path, "w", newline="") as csvfile:
                writer = csv.DictWriter(csvfile, fieldnames=tab_cols)
                if export_config["header"]:
                    writer.writeheader()
                writer.writerows(table_records)

    def export_split_records(self, table_split_record_dict: dict):
        for table_name, split_records in table_split_record_dict.items():
            table_col = split_records["__cols"]
            for schema, schema_records in split_records.items():
                if schema == "__cols":
                    continue
                tab_export_config = copy.deepcopy(self.base_config)
                tab_export_config["file_name"] = "%s.%s" % (table_name, tab_export_config["export_type"])
                tab_export_config["export_dir"] = os.path.join(tab_export_config["export_dir"], schema)
                if len(schema_records) > 0:
                    self.export_table_param(tab_export_config, schema_records, table_col)

